from problem import *
from osm_utils4 import *
from all import *
from time import time
from weightedAstar import *
import sys
from graph import *


class Move(ProblemAction):
	def __init__(self, start, target, cost):
		self.start = start;
		self.target = target;
		self.cost = cost;
		
	def __cmp__(self, other):
		if self.start != other.start:
			return self.start.__cmp__(other.start);
		return self.target.__cmp__(other.target);
	
	def __hash__(self):
		return hash(self.start + self.target);
	
	def __str__(self):
		return str(self.start) + "->" + str(self.target);

class Location(ProblemState):
	
	def __init__(self, identity, goal, costCalculator, countryMap, parent):
		self.identity = identity;
		self.goal = goal;
		self.costCalculator = costCalculator;
		self.countryMap = countryMap;
		self.parent = parent;
		
	def getSuccessors(self):
		successors = dict();
		myJunction = self.countryMap.GetJunction(self.identity);
		for l in myJunction.links:
			target = l.target;
			if target == self.parent:
				continue;
			cost = self.costCalculator.calculate(l);
			successors[Move(self.identity, target, cost)] = Location(target, self.goal, self.costCalculator, self.countryMap, self.identity);
		return successors;
	
	def isGoal(self):
		return (self.identity == self.goal);
	
	def __cmp__(self, other):
		return self.identity.__cmp__(other.identity);
	
	def __hash__(self):
		return hash(self.identity);
	
	def __str__(self):
		return "Location: " + str(self.identity);

#######################################################
# COST FUNCTIONS
#######################################################
class CostCalculator:
	def __init__(self, weightDistance, weightTime, weightPetrol, countryMap):
		self.weightDistance = weightDistance;
		self.weightTime = weightTime;
		self.weightPetrol = weightPetrol;
		self.countryMap = countryMap;
				
	def calculate(self, l):
		cost = 0;
		cost += self.weightDistance * l.distance;
		cost += self.weightTime *(float(l.distance) / float(l.speed));
		cost += self.weightPetrol * float(l.distance) / float(self.countryMap.PetrolConsumption(l.speed));
		return cost;

def distanceCostCalculator(countryMap):
	return CostCalculator(1,0,0, countryMap);

def timeCostCalculator(countryMap):
	return CostCalculator(0,1,0, countryMap);

def petrolCostCalculator(countryMap):
	return CostCalculator(0,0,1, countryMap);

#######################################################
# HEURISTICS
#######################################################

class weightedHeuristic:
	def __init__(self,  weightDistance, weightTime, weightPetrol, countryMap):
		self.weightDistance = weightDistance;
		self.weightTime = weightTime;
		self.weightPetrol = weightPetrol;
		self.countryMap = countryMap;
		speeds = [l.speed for j in self.countryMap.junctions for l in j.links ];
		self.maxSpeed = float(max(speeds));
		petrols = [self.countryMap.PetrolConsumption(speed) for speed in speeds];		
		self.maxPetrol = max(petrols)
	def evaluate(self, state):
		j1 = self.countryMap.GetJunction(state.identity);
		j2 = self.countryMap.GetJunction(state.goal);
		distance = float(self.countryMap.ComputeDistance(j1.lat, j1.lon, j2.lat, j2.lon));
		cost = 0;
		cost += self.weightDistance * distance;
		cost += self.weightTime *(distance / self.maxSpeed);
		cost += self.weightPetrol * distance / float(self.maxPetrol);
		return cost;
		
		
def distanceHeuristic(countryMap):
	return weightedHeuristic(1,0,0,countryMap);

def timeHeuristic(countryMap):
	return weightedHeuristic(0,1,0,countryMap);

def petrolHeuristic(countryMap):
	return weightedHeuristic(0,0,1,countryMap);


allUpdates = [];

class newAStar:
	def __init__(self, max_depth=infinity):
		pass
	
	def find(self, problem_state, heuristic):
		def evaluator(node):
			return node.path_cost + heuristic.evaluate(node.state)
	
		def queue_generator():
			return PriorityQueue(evaluator)
		
		open_states = PriorityQueue(evaluator)
		closed_states = {}
		
		used_links = [];
		m.ZeroChangeCounter();
		open_states.append(Node(problem_state))
		
		allUpdates = [];
		
		while open_states and len(open_states) > 0:
			
			updates = m.GetSpeedUpdates(used_links);
			if len(updates) > 0:
				allUpdates.insert(0, updates);
				print "updating now"
				
				# Applying updates to map
				for l, change in updates:
					u = l[0];
					v = l[1];
					j = m.GetJunction(u);
					for ll in j.links:
						if v == ll.target:
							ll.speed = int(ll.speed * (1+change));
							break;
							
				
				# Generating lists of nodes to throw
				sources = [update[0][0] for update in updates] 
				targets = [update[0][1] for update in updates]
				tothrow = []
				for item in [t[1] for t in open_states.q]:
					path = item.getPath()
					flag = False
					for n in path:
						if flag:
							tothrow.append(n.state)
							continue
						if n.state.identity in targets:
							tothrow.append(n.state)
							flag = True
				
				toThrowIdentities = [s.identity for s in tothrow];
				
				# Remove open states that we need to throw
				newQueue = PriorityQueue(evaluator);
				for item in [t[1] for t in open_states.q]:
					if item.state not in tothrow:
						 newQueue.append(item);
				open_states = newQueue;
				
				# Remove closed_states that we need to throw
				for item in tothrow:
					if item in closed_states: 
						del closed_states[item]
				
				#Putting some closed states as open
				for s in closed_states.keys():
					# if the node is the source of modified link
					if s.identity in sources:
						open_states.append(closed_states[s]);
						del closed_states[s]
					# if the node is the source of a link to a throwed node
					else:
						for t in [l.target for l in m.GetJunction(s.identity).links]:
							if t in toThrowIdentities:
								open_states.append(closed_states[s]);
								del closed_states[s]
								break;
			
			# Continuing as usual
			node = open_states.pop()
			
			if node.state.isGoal(): 
				return node.getPathActions()
			
			if (node.state not in closed_states) or (node.path_cost < closed_states[node.state].path_cost):
				used_links.extend([(node.state.identity, l.identity) for l in node.state.getSuccessors().values()]);		
				closed_states[node.state] = node
				open_states.extend(node.expand())
				
		return None  

class naiveAStar:
	def __init__(self):
		pass
		
	def find(self, problem_state, heuristic):
		def evaluator(node):
			return node.path_cost + heuristic.evaluate(node.state)
	
		def queue_generator():
			return PriorityQueue(evaluator)
		
		open_states = PriorityQueue(evaluator)
		closed_states = {}
		
		used_links = [];
		m.ZeroChangeCounter();
		open_states.append(Node(problem_state))
		
		while open_states and len(open_states) > 0:
			
			updates = m.GetSpeedUpdates(used_links);
			if len(updates) > 0 and len(allUpdates) > 0:
				print "updating now"
				
				updates = allUpdates.pop();
				
				# Applying updates to map
				for l, change in updates:
					u = l[0];
					v = l[1];
					j = m.GetJunction(u);
					for ll in j.links:
						if v == ll.target:
							ll.speed = int(ll.speed * (1+change));
							break;
				open_states = PriorityQueue(evaluator)
				closed_states = {}
				used_links = [];
				open_states.append(Node(problem_state))
			
			# Continuing as usual
			node = open_states.pop()
			
			if node.state.isGoal(): 
				return node.getPathActions()
			
			if (node.state not in closed_states) or (node.path_cost < closed_states[node.state].path_cost):
				used_links.extend([(node.state.identity, l.identity) for l in node.state.getSuccessors().values()]);		
				closed_states[node.state] = node
				open_states.extend(node.expand())
				
		return None

def run_problem(num, num2, start, goal, cost, m, h, s):
	startTime = time();
	a = s.find(Location(start, goal, cost, m, -1), h);
	totalTime = time() - startTime;
	
	if a == None:
		print num, num2, "start = ", start, "goal = ", goal, "not connected"
		return;
	
	j1 = m.GetJunction(start);
	j2 = m.GetJunction(goal);
	totalCost = sum([x.cost for x in a]);
	print num, num2, totalTime, m.ComputeDistance(j1.lat, j1.lon, j2.lat, j2.lon), len(a), totalCost;


def load_problems():
	f = open("problems.txt");
	problems = [];
	for l in f:
		l = l.split();
		start = int(l[0]);	
		goal = int(l[1]);
		problems.append((start, goal));
	return problems;

if __name__ == "__main__":	
	
	problems = load_problems();
		
	s1 = newAStar();
	s2 = naiveAStar();

	m = CountryMap();
	print "Loading map ..."	
	m.LoadMap2();
	print "Map loaded."

	m.car = "Toyota Rav4";	
	costCalc = timeCostCalculator(m);
	heuristicCalc = timeHeuristic(m);
	
	for p in range(10):
		start = problems[p][0];
		goal = problems[p][1];
		allUpdates = []
		run_problem(p, 1, start, goal, costCalc, m, heuristicCalc, s1);
		run_problem(p, 2, start, goal, costCalc, m, heuristicCalc, s2);
	
			